#pragma once

#include "common.hpp"
#include "log.hpp"
#include "win.hpp"

constexpr size_t BUFFER_SIZE = 8192;

class socket_server;

enum class socket_operation {
    SOCKET_RECV,         // receive data operation
    SOCKET_SEND,         // send data operation
    SOCKET_ACCEPT,       // accept connection operation
    NON_SOCKET_OPERATION // invalid operation
};

struct io_context : win::overlapped {
    win::socket socket;
    win::wsa_buf wsa_buffer;
    std::vector<char> buffer;
    std::function<void()> callback;

    socket_operation operation;

    io_context(const io_context&) = delete;
    io_context& operator=(const io_context&) = delete;
    io_context(io_context&&) = delete;
    io_context& operator=(io_context&&) = delete;

    io_context(win::socket socket, socket_operation operation, size_t default_buffer_size = 4096)
        : _OVERLAPPED()
        , socket(socket)
        , wsa_buffer{}
        , buffer(default_buffer_size, 0)
        , operation(operation) {
        clear_base();
        wsa_buffer.buf = buffer.data();
        wsa_buffer.len = static_cast<win::ulong>(buffer.size());
    }

    io_context(win::socket socket, socket_operation operation, std::vector<char> buf)
        : _OVERLAPPED()
        , socket(socket)
        , wsa_buffer{}
        , buffer(std::move(buf))
        , operation(operation) {
        clear_base();
        wsa_buffer.buf = buffer.data();
        wsa_buffer.len = static_cast<win::ulong>(buffer.size());
    }

    void clear_base() {
        memset(static_cast<win::overlapped*>(this), 0, sizeof(win::overlapped));
    }

    void clear_buffer() {
        std::fill(buffer.begin(), buffer.end(), static_cast<char>(0));
    }

    ~io_context() = default;
};

struct socket_context {
    socket_server& server_instance;

    win::socket socket{INVALID_SOCKET};
    win::sockaddr_in address;

    win::iocp_object* iocp_obj{nullptr};

    std::mutex m;
    std::vector<std::unique_ptr<io_context>> io_contexts;

    socket_context(socket_server& server_instance, win::socket socket)
        : server_instance(server_instance)
        , socket(socket)
        , address{}
        , iocp_obj(nullptr) {
        memset(&address, 0, sizeof(address));
    }

    socket_context(const socket_context&) = delete;
    socket_context& operator=(const socket_context&) = delete;
    socket_context(socket_context&&) = delete;
    socket_context& operator=(socket_context&&) = delete;
    ~socket_context();

    [[nodiscard]] io_context& create_io_context(socket_operation operation);
    [[nodiscard]] io_context& create_io_context(socket_operation, std::vector<char>);
    [[nodiscard]] bool remove_io_context(io_context* io_ctx);
};

class socket_server {
public:
    friend struct socket_context;

    using recv_hook_t = std::function<bool(win::socket, const win::sockaddr_in&, std::vector<char>)>;
    using accept_hook_t = std::function<void(win::socket, const win::sockaddr_in&)>;
    using send_hook_t = std::function<void(win::socket, const win::sockaddr_in&, size_t, size_t)>;
    using close_hook_t = std::function<void(win::socket, const win::sockaddr_in&)>;

    socket_server(u_short server_port = 80)
        : _server_port(server_port)
        , _server_ip("0.0.0.0") {
        _accept_count = std::thread::hardware_concurrency() + 2;
        _initialize_default_hooks();
    }

    void start();

    [[nodiscard]] bool send(win::socket sock, std::vector<char>, const std::function<void()>& callback = nullptr);

    void stop();

    bool close_socket(win::socket sock);

    void register_recv_hooks(const recv_hook_t& f);
    void register_accept_hooks(const accept_hook_t& f);
    void register_send_hooks(const send_hook_t& f);
    void register_close_hooks(const close_hook_t& f);

    [[nodiscard]] bool alive() const {
        return _server_context != nullptr;
    }

private:
    size_t _accept_count;

    std::mutex _connections_mutex;
    std::vector<std::unique_ptr<socket_context>> _client_contexts;

    std::unique_ptr<socket_context> _server_context;

    u_short _server_port;
    std::string _server_ip;

    LPFN_ACCEPTEX _acceptex{};
    LPFN_GETACCEPTEXSOCKADDRS _get_acceptex_sock_addrs{};

    accept_hook_t _accept_hook;
    recv_hook_t _recv_hook;
    send_hook_t _send_hook;
    close_hook_t _close_hook;

/**
 * @brief A wrapped error handler to print function name
 */
#define handle_error(SOCK_CONTEXT, IO_CONTEXT, ERROR_CODE) \
    _handle_error(SOCK_CONTEXT, IO_CONTEXT, ERROR_CODE,    \
        std::string(__FILE__) + ":" + std::to_string(__LINE__) + " " + nonstd::to_string(std::wstring(__FUNCTIONW__)))

#define handle_socket_error(SOCK_CONTEXT, IO_CONTEXT) handle_error(SOCK_CONTEXT, IO_CONTEXT, ::WSAGetLastError())

    void _handle_error(
        socket_context& socket_ctx, io_context& io_ctx, win::dword error_code, const std::string& source = "");

    void _initialize_default_hooks();

    socket_context* _get_socket_context(win::socket sock);

    static bool _test_socket_alive(win::socket sock) {
        int bytes_sent = ::send(sock, "", 0, 0);
        return bytes_sent != -1;
    }

    static void WINAPI worker_thread(_Inout_ win::tp_callback_instance* instance, _Inout_opt_ void* sock_context,
        _Inout_opt_ void* overlapped, _In_ win::ulong io_result, _In_ win::ulong_ptr bytes_trans,
        _Inout_ win::iocp_object* io);

    [[nodiscard]] socket_context& _create_new_client_context(win::socket client_socket);

    bool _remove_client_context(socket_context* socket_ctx, io_context* io_ctx);

    void _clear_contexts();

    static win::iocp_object* create_threadpool_io(win::socket sock, socket_context& socket_ctx) {
        return ::CreateThreadpoolIo(reinterpret_cast<win::handle>(sock), worker_thread, &socket_ctx, nullptr);
    }

    void _initialize_socket();

    void _post_accept(io_context& io_ctx);
    void _post_recv(socket_context& socket_ctx, io_context& io_ctx);
    void _post_send(socket_context& socket_ctx, io_context& io_ctx);

    void _accept(socket_context& socket_ctx, io_context& io_ctx, win::dword bytes_transferred);
    void _recv(socket_context& socket_ctx, io_context& io_ctx, win::dword bytes_transferred);
    void _send(socket_context& socket_ctx, io_context& io_ctx, win::dword bytes_transferred);
};